{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gen_path=r'D:\\coding\\melody-generator-gan\\src\\src\\save_\\22_03_29\\22_03_29_23_08_35\\gen_10.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "from comp.GeneratorRMC import GeneratorRMC\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import yaml\n",
    "from torch import optim\n",
    "import time\n",
    "\n",
    "from torch.utils import data\n",
    "from comp.DataloadernoConJP import DataloadernoConJP as Dataloader"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "torch.set_default_tensor_type(torch.DoubleTensor)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load lyrics_seq\n",
      "load discrete_attr_seq\n",
      "loaded data\n"
     ]
    }
   ],
   "source": [
    "class Gan6():\n",
    "    def __init__(self):\n",
    "        # super(Gan6, self).__init__()\n",
    "        self.save_per_epoch: int = 10\n",
    "        self.head_size: int = 256\n",
    "        self.num_heads: int = 2\n",
    "        self.mem_slots: int = 1\n",
    "        self.discriminator_input_dim_music: int = 3\n",
    "        self.lyrics_dis_rate: float = 0.5\n",
    "        self.train_data_iterator = None\n",
    "        self.sequence_len: int = 10\n",
    "        self.train_D_steps: int = 1\n",
    "        self.train_G_steps: int = 1\n",
    "        self.total_epoch: int = 30\n",
    "        self.this_time = time.strftime(\"%y_%m_%d_%H_%M_%S\", time.localtime())\n",
    "        self.this_day = time.strftime(\"%y_%m_%d\", time.localtime())\n",
    "        # default config, you should change in yaml file.\n",
    "        self.data_params: dict = {'batch_size': 100,\n",
    "                                  'shuffle': True,\n",
    "                                  'num_workers': 6}\n",
    "        # This dic to load data\n",
    "        self.load_flag: bool = False\n",
    "        self.learning_rate_D: float = 0.1\n",
    "        self.learning_rate_G: float = 0.0001\n",
    "        self.discriminator_out_dim: int = 1\n",
    "        self.discriminator_input_dim: int = 1\n",
    "        self.generator_out_dim: int = 3\n",
    "        self.hidden_dim: int = 400\n",
    "        self.ff1_out: int = 400\n",
    "        self.lyrics_dim: int = 64\n",
    "        self.embed_dim: int = 32\n",
    "        self.cuda: bool = False\n",
    "        self.read_yaml('D:\\coding\\melody-generator-gan\\src\\config.yaml')\n",
    "\n",
    "        self.device = torch.device('cuda:0' if self.cuda else 'cpu')\n",
    "\n",
    "        # self.hidden_dim_for_rmc = self.mem_slots * self.num_heads * self.head_size\n",
    "\n",
    "        self.gen = GeneratorRMC(mem_slots=self.mem_slots, num_heads=self.num_heads, head_size=self.head_size,\n",
    "                                embed_dim=self.lyrics_dim, ff1_out=self.ff1_out,\n",
    "                                hidden_dim_lstm_as_input_size=self.embed_dim, out_dim=self.generator_out_dim,\n",
    "                                cuda=self.cuda, init_batch_size=self.data_params['batch_size'],sentence_len=self.sequence_len)\n",
    "        self.gen.load_state_dict(torch.load(gen_path))\n",
    "\n",
    "\n",
    "        self.load_data(from_pth=False)\n",
    "\n",
    "        if self.cuda:\n",
    "            self.gen.to(self.device)\n",
    "\n",
    "    def read_yaml(self, path):\n",
    "        with open(path, 'r') as config_yaml:\n",
    "            data = yaml.load(config_yaml, yaml.FullLoader)\n",
    "        self.cuda = data['cuda']\n",
    "        self.embed_dim = data['embed_dim']\n",
    "        self.lyrics_dim = 2 * self.embed_dim\n",
    "        self.ff1_out = data['ff1_out']\n",
    "        self.hidden_dim = data['hidden_dim']\n",
    "        self.generator_out_dim = data['generator_out_dim']\n",
    "        self.discriminator_input_dim = self.embed_dim + self.generator_out_dim\n",
    "        self.discriminator_input_dim_music = self.generator_out_dim\n",
    "        self.discriminator_out_dim = data['discriminator_out_dim']\n",
    "        self.learning_rate_G = data['learning_rate_G']\n",
    "        self.learning_rate_D = data['learning_rate_D']\n",
    "        self.load_flag = data['load_model']\n",
    "        self.data_params = data['data_params']\n",
    "        self.total_epoch = data['epochs']\n",
    "        self.train_G_steps = data['train_G_steps']\n",
    "        self.train_D_steps = data['train_D_steps']\n",
    "        self.sequence_len = data['sequence_len']\n",
    "        self.lyrics_dis_rate = data['lyrics_dis_rate']\n",
    "        self.mem_slots = data['mem_slots']\n",
    "        self.num_heads = data['num_heads']\n",
    "        self.head_size = data['head_size']\n",
    "        self.save_per_epoch = data['save_per_epoch']\n",
    "\n",
    "\n",
    "    def load_data(self, from_pth=False):\n",
    "        if not from_pth:\n",
    "            training_set = Dataloader('2022-01-15_18_30_13_embeddings10_vector.pt',\n",
    "                                      '2022-01-15_18_30_13_vocabulary10_lookup.json',\n",
    "                                      self.sequence_len)\n",
    "            self.train_data_iterator = data.DataLoader(training_set, **self.data_params)\n",
    "            print('loaded data')\n",
    "        else:\n",
    "            self.train_data_iterator = torch.load('../teswt_set.pth')\n",
    "            print('loaded data')\n",
    "\n",
    "    def get_gen_data(self):\n",
    "        epoch=999\n",
    "        self.gen.eval()\n",
    "\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for i, test_data in enumerate(self.train_data_iterator):\n",
    "                lyrics_seq = test_data[0].to(self.device)\n",
    "\n",
    "                discrete_val_seq = test_data[1].to(self.device)\n",
    "\n",
    "                gen_out = self.gen(lyrics_seq)\n",
    "                base_data_root = os.path.join(os.getcwd(), 'src/save_/{}/{}/'.format(self.this_day, self.this_time))\n",
    "                if not os.path.exists(base_data_root):\n",
    "                    os.makedirs(base_data_root)\n",
    "                np.save(os.path.join(base_data_root, 'epoch{}_gen_data_{}'.format(epoch, i)),\n",
    "                        torch.cat((discrete_val_seq, gen_out), dim=2))\n",
    "\n",
    "\n",
    "\n",
    "                # generator.zero_grad()\n",
    "                # discriminator.zero_grad()\n",
    "\n",
    "                # gen_out = self.gen(lyrics_seq, noise_seq)\n",
    "                # print(gen_out[2])\n",
    "                #\n",
    "                # dis_out = self.dis(gen_out, lyrics_seq)\n",
    "                #\n",
    "                # print('generator out loss : {} ,discriminator out loss : {}'.format(gen_out, dis_out))\n",
    "                if i == 5:\n",
    "                    break\n",
    "\n",
    "gen6=Gan6()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "outputs": [],
   "source": [
    "gen6.get_gen_data()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}